package tsj.ai.alibaba.graph.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-06 00:18:22
 */
public class ExpanderNode  implements NodeAction {

    // 扩展文本的提示词
    private static String template = "You are an expert at information retrieval and search optimization.\n" +
            "Your task is to generate {number} different versions of the given query." +
            "Each variant must cover different perspectives or aspects of the topic," +
            "while maintaining the core intent of the original query. The goal is to" +
            "expand the search space and improve the chances of finding relevant information." +
            "Do not explain your choices or add any other text." +
            "Provide the query variants separated by newlines." +
            "Original query: {query}" +
            "Query variants:";
    private static final PromptTemplate DEFAULTPROMPTTEMPLATE = new PromptTemplate(template);

    // 调用 AI 模型的 client 端
    private final ChatClient chatClient;

    // 默认扩展为 3 条相似问题
    private final Integer NUMBER = 3;

    public ExpanderNode(ChatClient.Builder chatClientBuilder) {
        this.chatClient = chatClientBuilder.build();
    }

    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        String query = state.value("query", "");
        Integer expanderNumber = state.value("expandernumber", this.NUMBER);


        Flux<String> streamResult = this.chatClient.prompt()
                .user((user) -> user.text(DEFAULTPROMPTTEMPLATE.getTemplate())
                .param("number", expanderNumber).param("query", query))
                .stream().content();
        String result = streamResult.reduce("", (acc, item) -> acc + item).block();
        List<String> queryVariants = Arrays.asList(result.split("\n"));

        HashMap<String, Object> resultMap = new HashMap<>();
        resultMap.put("expandercontent", queryVariants);
        return resultMap;
    }
}
